#include <iostream>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <thread>
#include <fstream>
#include <vector>
#include <chrono>
#include <sstream>
#ifdef _WIN32
#pragma comment(lib, "ws2_32.lib")
#endif
using namespace std;

// 日志函数
void LogMessage(const string& level, const string& message) {
    auto now = chrono::system_clock::now();
    time_t now_time = chrono::system_clock::to_time_t(now);
    tm tm_info;
    localtime_s(&tm_info, &now_time);
    char time_str[20];
    strftime(time_str, sizeof(time_str), "%Y-%m-%d %H:%M:%S", &tm_info);

    ofstream logfile("network.log", ios::app);
    logfile << "[" << time_str << "][" << level << "] " << message << endl;
    cout << "[" << time_str << "][" << level << "] " << message << endl;
}


// Winsock初始化
bool InitializeWinsock() {
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2,2), &wsaData) != 0) {
        LogMessage("FATAL", "WSAStartup failed: " + to_string(WSAGetLastError()));
        return false;
    }
    if (LOBYTE(wsaData.wVersion) != 2 || HIBYTE(wsaData.wVersion) != 2) {
        LogMessage("ERROR", "Winsock version 2.2 not available");
        WSACleanup();
        return false;
    }
    return true;
}


const int BUF_SIZE = 64;
SOCKET sServer;
SOCKET sClient;
WSADATA wsd;
SOCKADDR_IN servAddr;
char bufRecv[BUF_SIZE];
char sendBuf[BUF_SIZE];
int retVal;

// 清理函数
void CleanupWinsock() {
    WSACleanup();
}



// UDP广播线程函数
void BroadcastServer(int tcpPort) {
    SOCKET udpSocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (udpSocket == INVALID_SOCKET) {
        LogMessage("ERROR", "创建广播套接字失败");
        return;
    }

    // 设置广播选项
    int broadcast = 1;
    if (setsockopt(udpSocket, SOL_SOCKET, SO_BROADCAST, (char*)&broadcast, sizeof(broadcast)) == SOCKET_ERROR) {
        LogMessage("ERROR", "设置广播选项失败");
        closesocket(udpSocket);
        return;
    }

    SOCKADDR_IN broadcastAddr;
    broadcastAddr.sin_family = AF_INET;
    broadcastAddr.sin_port = htons(4998); // 广播端口
    inet_pton(AF_INET, "255.255.255.255", &broadcastAddr.sin_addr);

    while (true) {
        string msg = "SERVER_PORT:" + to_string(tcpPort);
        if (sendto(udpSocket, msg.c_str(), static_cast<int>(msg.size()), 0, (sockaddr*)&broadcastAddr, sizeof(broadcastAddr)) == SOCKET_ERROR) {
            LogMessage("ERROR", "广播发送失败");
            break;
        }
        LogMessage("INFO", "正在广播服务端端口: " + to_string(tcpPort));
        this_thread::sleep_for(chrono::seconds(3));
    }
    closesocket(udpSocket);
}

// 服务端功能
int TestServer(int port) {
    // 初始化地址结构
    servAddr.sin_family = AF_INET;
    servAddr.sin_port = htons(port);
    servAddr.sin_addr.s_addr = INADDR_ANY;

    // 绑定套接字
    if (bind(sServer, (LPSOCKADDR)&servAddr, sizeof(SOCKADDR_IN))) {
        LogMessage("ERROR", "套接字绑定失败");
        closesocket(sServer);
        WSACleanup();
        return -1;
    }
    LogMessage("INFO", "服务端已绑定端口: " + to_string(port));

    // 开始监听
    if (listen(sServer, 1) == SOCKET_ERROR) {
        LogMessage("ERROR", "监听失败");
        closesocket(sServer);
        WSACleanup();
        return -1;
    }
    LogMessage("INFO", "服务端开始监听...");

    // 启动广播线程
    thread broadcastThread(BroadcastServer, port);
    broadcastThread.detach();

    // 接受客户端连接
    sockaddr_in clientAddr;
    int addrLen = sizeof(clientAddr);
    sClient = accept(sServer, (sockaddr*)&clientAddr, &addrLen);
    if (sClient == INVALID_SOCKET) {
        LogMessage("ERROR", "接受连接失败");
        closesocket(sServer);
        WSACleanup();
        return -1;
    }

    char clientIP[INET_ADDRSTRLEN];
    inet_ntop(AF_INET, &clientAddr.sin_addr, clientIP, INET_ADDRSTRLEN);
    LogMessage("INFO", "客户端已连接: " + string(clientIP) + ":" + to_string(ntohs(clientAddr.sin_port)));

    // 服务端接收线程函数
    thread recvThread([&]() {
        char threadBuf[BUF_SIZE];
        while (true) {
            ZeroMemory(threadBuf, BUF_SIZE);
            int ret = recv(sClient, threadBuf, BUF_SIZE, 0);
            if (ret <= 0) {
                LogMessage("WARN", "客户端已断开");
                break;
            }
            LogMessage("RECV", "来自客户端: " + string(threadBuf));
            cout << "\r\033[K\033[34m[客户端说]\033[0m " << threadBuf << "\n\033[33m服务端发送: \033[0m";
            cout << "服务端发送: ";
            fflush(stdout);
        }
    });

    // 发送循环
    while (true) {
        cout << "服务端发送: ";
        cin.getline(sendBuf, BUF_SIZE);
        
        if (sendBuf[0] == '0') {
            cout << "确定要关闭服务器吗？(y/n): ";
            char confirm;
            cin >> confirm;
            cin.ignore();
            if (confirm == 'y' || confirm == 'Y') break;
            continue;
        }

        send(sClient, sendBuf, static_cast<int>(strlen(sendBuf)), 0);
        LogMessage("SEND", "发送给客户端: " + string(sendBuf));
    }
    
    recvThread.detach();

    if (sClient != INVALID_SOCKET) closesocket(sClient);
    if (sServer != INVALID_SOCKET) closesocket(sServer);
    WSACleanup();
        cout << "\n\n输入任意字符返回主菜单...";
    cin.get();
    cin.ignore();
    return 0;
}

// 客户端发现服务端
vector<pair<string, int>> DiscoverServers() {
    vector<pair<string, int>> servers;
    SOCKET udpSocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (udpSocket == INVALID_SOCKET) return servers;

    // 设置接收超时
    DWORD timeout = 3000;
    setsockopt(udpSocket, SOL_SOCKET, SO_RCVTIMEO, (char*)&timeout, sizeof(timeout));

    sockaddr_in recvAddr;
    recvAddr.sin_family = AF_INET;
    recvAddr.sin_port = htons(4998);
    recvAddr.sin_addr.s_addr = INADDR_ANY;
    bind(udpSocket, (sockaddr*)&recvAddr, sizeof(recvAddr));

    char buffer[1024];
    sockaddr_in serverAddr;
    int addrLen = sizeof(serverAddr);

    LogMessage("INFO", "正在搜索可用服务端...");
    while (true) {
        int bytes = recvfrom(udpSocket, buffer, sizeof(buffer), 0, (sockaddr*)&serverAddr, &addrLen);
        if (bytes == SOCKET_ERROR) break;

        string msg(buffer, bytes);
        if (msg.find("SERVER_PORT:") == 0) {
            int port = stoi(msg.substr(12));
            char ip[INET_ADDRSTRLEN];
            inet_ntop(AF_INET, &serverAddr.sin_addr, ip, INET_ADDRSTRLEN);
            servers.emplace_back(ip, port);
            LogMessage("INFO", "发现服务端: " + string(ip) + ":" + to_string(port));
        }
    }
    closesocket(udpSocket);
    return servers;
}

// 客户端功能
int TestClient(const string& ip, int port) {
    // 设置服务器地址
    servAddr.sin_family = AF_INET;
    servAddr.sin_port = htons(port);
    inet_pton(AF_INET, ip.c_str(), &servAddr.sin_addr);

    // 连接服务器
    if (connect(sServer, (sockaddr*)&servAddr, sizeof(servAddr))) {
        LogMessage("ERROR", "连接服务器失败");
        closesocket(sServer);
        WSACleanup();
        return -1;
    }
    LogMessage("INFO", "已连接到服务器: " + ip + ":" + to_string(port));

    // 客户端接收线程函数
    thread recvThread([&]() {
        char threadBuf[BUF_SIZE];
        while (true) {
            ZeroMemory(threadBuf, BUF_SIZE);
            int ret = recv(sServer, threadBuf, BUF_SIZE, 0);
            if (ret <= 0) {
                LogMessage("WARN", "连接已断开");
                break;
            }
            LogMessage("RECV", "来自服务端: " + string(threadBuf));
            cout << "\r\033[K\033[32m[服务端说]\033[0m " << threadBuf << "\n\033[33m客户端发送: \033[0m";
            cout << "客户端发送: ";
            fflush(stdout);
        }
    });

    // 发送循环
    while (true) {
        cout << "客户端发送: ";
        cin.getline(sendBuf, BUF_SIZE);
        
        if (sendBuf[0] == '0') {
            cout << "确定要退出吗？(y/n): ";
            char confirm;
            cin >> confirm;
            cin.ignore();
            if (confirm == 'y' || confirm == 'Y') break;
            continue;
        }

        send(sServer, sendBuf, static_cast<int>(strlen(sendBuf)), 0);
        LogMessage("SEND", "发送给服务端: " + string(sendBuf));
    }
    
    recvThread.detach();

    closesocket(sServer);
    WSACleanup();
        cout << "\n\n输入任意字符返回主菜单...";
    cin.get();
    cin.ignore();
    return 0;
}

int main() {
    // 设置控制台编码为UTF-8
    SetConsoleOutputCP(CP_UTF8);
    SetConsoleCP(CP_UTF8);
    if (!InitializeWinsock()) {
        return 1;
    }
    atexit(CleanupWinsock);

    // 初始化Winsock
    if (WSAStartup(MAKEWORD(2, 2), &wsd)) {
        LogMessage("ERROR", "Winsock初始化失败");
        return -1;
    }

    // 创建主套接字
    sServer = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (sServer == INVALID_SOCKET) {
        LogMessage("ERROR", "套接字创建失败");
        WSACleanup();
        return -1;
    }

    // 主菜单
    int choice;
    cout << "\033[36m╔═══════════════╗\n║\033[0m  \033[33m网络通信终端\033[0m  \033[36m║\n╚═══════════════╝\033[0m\n";
cout << "\033[36m1.\033[0m \033[32m启动服务端\033[0m\n";
cout << "\033[36m2.\033[0m \033[32m启动客户端\033[0m\n";
cout << "\033[33m请选择操作 > \033[0m";
    cin >> choice;
    cin.ignore();

    if (choice == 1) {
        int port;
        cout << "输入服务端端口 (默认4999): ";
        string portInput;
        getline(cin, portInput);
        port = portInput.empty() ? 4999 : stoi(portInput);
        while (true) {
            TestServer(port);
            LogMessage("WARN", "服务端连接已断开，等待5秒后重启...");
            this_thread::sleep_for(chrono::seconds(5));
        }
    } else if (choice == 2) {
        cout << "\n\033[35m──┤ 连接方式 ├──\033[0m\n";
cout << "\033[36m1.\033[0m \033[32m手动输入\033[0m\n";
cout << "\033[36m2.\033[0m \032m自动搜索\033[0m\n";
cout << "\033[33m请选择 > \033[0m";
        int mode;
        cin >> mode;
        cin.ignore();

        if (mode == 1) {
            string ip;
            int port;
            cout << "输入服务器IP: ";
            getline(cin, ip);
            cout << "输入服务器端口: ";
            cin >> port;
            cin.ignore();
            TestClient(ip, port);
        } else if (mode == 2) {
            auto servers = DiscoverServers();
            while (servers.empty()) {
                LogMessage("WARN", "未找到可用服务端，5秒后重试...");
                this_thread::sleep_for(chrono::seconds(5));
                servers = DiscoverServers();
            }

            cout << "\n\033[35m──┤ 可用服务端列表 ├──\033[0m\n";
            for (size_t i = 0; i < servers.size(); ++i) {
                cout << i + 1 << ". " << servers[i].first << ":" << servers[i].second << endl;
            }

            int choice;
            cout << "选择要连接的服务端: ";
            cin >> choice;
            cin.ignore();

            if (choice > 0 && choice <= servers.size()) {
                TestClient(servers[choice - 1].first, servers[choice - 1].second);
            } else {
                LogMessage("ERROR", "无效选择");
            }
        }
    }

        cout << "\n\n输入任意字符返回主菜单...";
    cin.get();
    cin.ignore();
    return 0;
}